import numpy as np
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns


def get_result(prediction):
    out = []
    for i in range(len(prediction)):
        out.append(np.argmax(prediction[i]))
        # print(i,out,prediction[i])
    return np.array(out)


label = np.load("../model_file/pre_result/label_4.npy")
pred = np.load("../model_file/pre_result/pred_4.npy")
prediction = get_result(pred)

fig = plt.figure(figsize=(8,6),dpi=300)
ax = fig.add_subplot(111)
conf_mat = confusion_matrix(label, prediction)
sns.set(font_scale=1.2)
sns.heatmap(conf_mat,
            annot=True,
            xticklabels=['AS', 'MS', 'MR', 'MVP'],
            yticklabels=['AS', 'MS', 'MR', 'MVP'],
            cmap = plt.cm.YlGnBu
            )
ax.set_title('Confusion Matrix(Four Categories)',fontsize = 20)
ax.set_xlabel('Prediction Label',fontsize = 16)
ax.set_ylabel('True Label',fontsize = 16)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.show()
